# -*- coding:utf-8 -*-

# 使用 Python 生成类似于下图中的字母+数字验证码图片
import numpy
from PIL import Image, ImageDraw, ImageFont, ImageFilter

import RotateImg
import utils
from utils import WriteXml, FindNumberEdge, ReviseInfos
from tqdm import tqdm
import numpy as np
import random
import mycv
import cv2
import os
import string


# 随机获取四个字母
def get_random_char(length, NumStart):
    # chars = string.ascii_letters + string.digits
    # chars = string.digits
    # chars = [str(x) for x in range(10, 131, 5)]
    chars = ["10", "11", "12", "13", "15"]
    chars_candidate = []
    for i, Num in enumerate(NumStart):
        if Num < NumEnd:
            chars_candidate.append(chars[i])
    chars_out = [random.choice(chars_candidate) for _ in range(length)]
    for i in range(length):
        NumStart[chars.index(chars_out[i])] += 1
    return chars_out, NumStart


class GetColor:

    # 获取颜色
    @staticmethod
    def DefaultRandom():
        return random.randint(30, 100), random.randint(30, 100), random.randint(30, 100)

    @staticmethod
    def White():
        return random.randint(200, 255), random.randint(200, 255), random.randint(200, 255)

    @staticmethod
    def WhiteFully():
        return 255, 255, 255

    @staticmethod
    def WhiteGray():
        hsv = np.zeros((1, 1, 3), dtype=np.uint8)
        if random.random() < 0.8:
            hsv[0, 0, 0] = random.randint(0, 10)
        else:
            hsv[0, 0, 0] = random.randint(156, 180)
        hsv[0, 0, 1] = random.randint(0, 20)
        hsv[0, 0, 2] = random.randint(100, 255)
        rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
        return rgb[0, 0, 0], rgb[0, 0, 1], rgb[0, 0, 2]

    @staticmethod
    def Black():
        hsv = np.zeros((1, 1, 3), dtype=np.uint8)
        hsv[0, 0, 0] = random.randint(0, 180)
        hsv[0, 0, 1] = random.randint(0, 255)
        hsv[0, 0, 2] = random.randint(0, 46)
        rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
        return int(rgb[0, 0, 0]), int(rgb[0, 0, 1]), int(rgb[0, 0, 2])

    @staticmethod
    def BlackFully():
        return 0, 0, 0

    @staticmethod
    def Color180():
        return 180, 180, 180

    @staticmethod
    def RandColor(LowerBound, UpperBound):
        return random.randint(LowerBound, UpperBound), random.randint(LowerBound, UpperBound), random.randint(
            LowerBound, UpperBound)

    @staticmethod
    def Random0255():
        return random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)


# 获取验证码图片
def get_codepic(OutPath, PathNoiseBig, ImgNum, ImgSize, NumberSize, NameOffset, NumStart, NumEnd, col_nums,
                row_nums, flags):
    Color = GetColor
    ttf_list, NoiseBigList = [], []
    for _, _, files in os.walk(r"fonts"):
        for file in files:
            ttf_list.append(file)

    for _, _, files in os.walk(PathNoiseBig):
        for file in files:
            NoiseBigList.append(file)

    if NumEnd:
        ImgNum = 0
        for x in NumStart:
            ImgNum += NumEnd - x
        ImgNum = int(ImgNum / col_nums / row_nums)
    if flags[0]:
        col2 = col_nums + 3
    else:
        col2 = col_nums
    for NameNum in tqdm(range(ImgNum)):

        scale = random.randint(2, 3)

        BackGroundColor = Color.Random0255()

        width, height = random.randint(int(ImgSize[0] * 0.8), int(ImgSize[0] * 1.25)) * scale, random.randint(
            int(ImgSize[1] * 0.8), int(ImgSize[1] * 1.25)) * scale
        image = (np.ones((height, width, 3)) * BackGroundColor).astype(np.uint8)

        # 创建验证码图片
        code, NumStart = get_random_char(col_nums * row_nums, NumStart)

        Infos = [OutPath + "/N%05d.xml" % (NameNum + NameOffset)]

        for code_index in range(len(code)):
            FontSize = int(min(int(width / col2), int(height / row_nums)) * 0.6 * random.uniform(0.5, 1))
            # FontSize = 25 * scale

            font = ImageFont.truetype('./fonts/' + random.choice(ttf_list), size=FontSize)
            img_text_size = int(FontSize * 2.5)
            img_text2 = Image.new('RGB', (img_text_size, img_text_size),
                                  (BackGroundColor[2], BackGroundColor[1], BackGroundColor[0]))
            draw = ImageDraw.Draw(img_text2)
            color_text = mycv.get_different_color(BackGroundColor, 70)
            draw.text((int((img_text_size - FontSize) / 2), int((img_text_size - FontSize) / 2)), code[code_index], font=font, fill=color_text)
            img_text2 = cv2.cvtColor(np.asarray(img_text2), cv2.COLOR_RGB2BGR)

            img_text2 = RotateImg.rotate_img(img_text2, [-30, 30], BackGroundColor)

            XMin, YMin, XMax, YMax = FindNumberEdge(img_text2, BackGroundColor)

            # try:
            RandH = random.randint(0, int(height / row_nums - YMax + YMin))
            RandW = random.randint(0, int(width / col2 - XMax + XMin))
            # except:
            #     print("ff")

            y_start = int(int(code_index / col_nums) * height / row_nums) + RandH
            y_end = int(int(code_index / col_nums) * height / row_nums) + RandH + YMax - YMin
            x_start = int((code_index % col_nums) * width / col2) + RandW
            x_end = int((code_index % col_nums) * width / col2) + RandW + XMax - XMin
            image[y_start:y_end, x_start:x_end] = img_text2[YMin:YMax, XMin:XMax]

            Infos.append(
                [x_start, y_start, x_end, y_end, code[code_index]])

        if flags[0]:
            image, Infos[1:] = utils.ResizeJpgAndXml(image, Infos[1:], 600, 300)
            x_end = int(col_nums * image.shape[1] / col2)
            NoisePic1 = cv2.resize(cv2.imread(PathNoiseBig + '/' + random.choice(NoiseBigList)), (100, 100))
            NoisePic2 = cv2.resize(cv2.imread(PathNoiseBig + '/' + random.choice(NoiseBigList)), (100, 100))
            NoisePic3 = cv2.resize(cv2.imread(PathNoiseBig + '/' + random.choice(NoiseBigList)), (100, 100))
            NoisePic4 = cv2.resize(cv2.imread(PathNoiseBig + '/' + random.choice(NoiseBigList)), (100, 100))
            NoisePic = np.hstack((np.vstack((NoisePic1, NoisePic2)), (np.vstack((NoisePic3, NoisePic4)))))

            H_noise = np.random.randint(int(image.shape[0] / 4), int(image.shape[0] / 2))
            W_noise = np.random.randint(int((image.shape[1] - x_end) / 4), int(image.shape[1] - x_end))
            NoisePic = cv2.resize(NoisePic, (W_noise, H_noise))
            NoisePic = RotateImg.rotate_img(NoisePic, [-5, 5])
            RandH = random.randint(0, image.shape[0] - H_noise)
            RandW = random.randint(0, int(image.shape[1] - x_end) - W_noise) + x_end
            image[RandH:RandH + H_noise, RandW:RandW + W_noise] = NoisePic

            # 加直线
            if flags[1]:
                mycv.draw_line(image, 20, Color.Black(), 2)

            # 填充噪点
            noise = np.random.randint(0, 256, size=image.shape, dtype=np.uint8)
            image = cv2.addWeighted(image, 0.9, noise, 0.1, 0)

            # 模糊处理
            image = cv2.GaussianBlur(image, (5, 5), 2, 1)
        else:

            # 加直线
            if flags[1]:
               mycv.draw_line(image, 7, Color.Black(), random.randint(1, 2))

            # 填充噪点
            noise = np.random.randint(0, 256, size=image.shape, dtype=np.uint8)
            image = cv2.addWeighted(image, 0.9, noise, 0.1, 0)

            # 模糊处理
            if random.random() < 0.5:
                image = cv2.GaussianBlur(image, (5, 5), 2, 1)
            image = mycv.motion_blur(image, random.randint(1, 6))

            image, Infos[1:] = utils.ResizeJpgAndXml(image, Infos[1:], 600, 300)

        cv2.imwrite(Infos[0][:-3] + "jpg", image)
        WriteXml(Infos, image.shape[1], image.shape[0])


if __name__ == '__main__':
    OutPath = r"/media/fang/TOSHIBA EXT/Temp/600X300/N_11"
    NoiseBig = r'/media/fang/TOSHIBA EXT/BackGroundpic/noise/Big'
    ImgNum = 0
    NameOffset = 0
    row_numbers = 1
    ImgSize = (140, 70)
    # NumberSize = (110, 130)
    # NumberSize = (12, 24)
    NumberSize = (30, 60)
    col_numbers = int(ImgSize[0] / 2 / NumberSize[0])
    flag_noise = False
    flag_line = True

    flags = [flag_noise, flag_line]



    # NumStart = [10661, 3370, 1540, 1637, 1252, 3958, 1518, 1497, 1696, 286, 0]
    # NumCircle = [0, 384, 898, 161, 910, 103, 1142, 110, 1325, 260, 1291, 227, 1233, 264, 1466, 230, 81, 205, 299, 166, 445, 263, 481, 0, 624]
    # NumRect =   [251, 472, 260, 300, 220, 274, 237, 607, 242, 220, 226, 258, 255, 262, 261, 258, 239, 267, 0, 0, 0, 0, 0, 0, 0]
    # NumStart = [NumCircle[idx] + NumRect[idx] for idx in range(len(NumCircle))]
    # NumStart = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    NumStart = [930, 1416, 962, 1248, 768]
    NumEnd = 25000
    if flag_noise:
        col_numbers -= 3
    get_codepic(OutPath, NoiseBig, ImgNum, ImgSize, NumberSize, NameOffset, NumStart, NumEnd, col_numbers, row_numbers, flags)
